import backtrader as bt
import akshare as ak
import pandas as pd

class StockScreener(bt.Strategy):
    params = (
        ('period', 20),
    )

    def __init__(self):
        self.sma = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.period)
        self.selected_stocks = set()

    def next(self):
        if self.data.close[0] > self.sma[0]:
            self.selected_stocks.add(self.data._name)

def fetch_data(stock_code, start_date, end_date):
    stock_data = ak.stock_zh_a_hist(symbol=stock_code, start_date=start_date, end_date=end_date, adjust="qfq")
    stock_data.index = pd.to_datetime(stock_data['日期'])
    stock_data = stock_data[['开盘', '最高', '最低', '收盘', '成交量']]
    stock_data.columns = ['open', 'high', 'low', 'close', 'volume']
    return stock_data

def run_screener(stock_codes, start_date, end_date):
    cerebro = bt.Cerebro()

    for stock_code in stock_codes:
        data = bt.feeds.PandasData(dataname=fetch_data(stock_code, start_date, end_date))
        cerebro.adddata(data, name=stock_code)

    cerebro.addstrategy(StockScreener)

    results = cerebro.run()
    return results[0].selected_stocks

if __name__ == '__main__':
    stock_codes = ['000001', '000002', '000004']  # 示例股票代码
    start_date = '20200101'
    end_date = '20210101'

    selected_stocks = run_screener(stock_codes, start_date, end_date)
    print("Selected Stocks:", selected_stocks)